import csv


def generate_arithmetic_problem():
    import random
    operators = ['+', '-', '*', '**']
    op_templates = {
        '+': [
            lambda x, y: f"{x} + {y} = ",
            lambda x, y: f"{x}+{y}=",
            lambda x, y: f"{x} plus {y} equals ",
            lambda x, y: f"{x}加{y}等于",
            lambda x, y: f"{x}+{y}等于",
        ],
        '-': [
            lambda x, y: f"{x} - {y} = ",
            lambda x, y: f"{x}-{y}=",
            lambda x, y: f"{x} minus {y} equals ",
            lambda x, y: f"{x}减{y}等于",
            lambda x, y: f"{x}-{y}等于",
        ],
        '*': [
            lambda x, y: f"{x} * {y} = ",
            lambda x, y: f"{x}*{y}=",
            lambda x, y: f"{x} times {y} equals ",
            lambda x, y: f"{x}乘{y}等于",
            lambda x, y: f"{x}*{y}等于",
        ],
        '**': [
            lambda x, y: f"{x} ** {y} = ",
            lambda x, y: f"{x}**{y}=",
            lambda x, y: f"{x} to the power of {y} equals ",
            lambda x, y: f"{x}的{y}次方等于",
            lambda x, y: f"{x}的{y}次幂等于",
            lambda x, y: f"{x}**{y}等于",
        ]
    }
    operator = random.choice(operators)
    templates = op_templates[operator]
    template = random.choice(templates)
    loperand = random.randint(1, 100)
    if operator in ['+', '*']:
        roperand = random.randint(1, 100)
    elif operator == '-':
        roperand = random.randint(1, loperand)
    elif operator == '**':
        roperand = random.randint(1, 5)
    else: raise Exception(f"Unknown operator: {operator}")
    expression = f"{loperand} {operator} {roperand}"
    answer = eval(expression)
    return template(loperand, roperand), operator, answer


if __name__ == '__main__':
    file_path = 'dataset/test.csv'

    header = ['problem', 'operator', 'answer']
    # data = []

    with open(file_path, mode='w', newline='', encoding='utf-8') as file:
        writer = csv.writer(file)
        writer.writerow(header)
        
        for row in range(10000):
            problem, op, answer = generate_arithmetic_problem()
            writer.writerow([problem, op, str(answer)])

    print(f'CSV 文件已生成：{file_path}')
